# -*- coding: utf-8 -*-
"""
Created on Wed Aug 26 15:47:14 2020

@author: litao
"""

# -*- coding: utf-8 -*-
"""
Created on Fri Oct 18 19:59:44 2019

@author: litao
"""
import os.path as op

from mne.filter import next_fast_len

import mne


import scipy.io as sio
#subject_no = 2
subject_names = ['sunxiangyu_20140511.mat','wangkui_20140620.mat','dujingcheng_20131027.mat','xiayulu_20140527.mat','liuqiujun_20140621.mat']
abbreviation = ['sxy','wk','djc','xyl','lqj']

channels_name = ['FP1','FPZ','FP2','AF3','AF4','F7','F5','F3','F1','FZ','F2','F4','F6','F8','FT7','FC5','FC3','FC1','FCZ','FC2','FC4','FC6','FT8','T7','C5','C3','C1','CZ','C2','C4','C6','T8','TP7','CP5','CP3','CP1','CPZ','CP2','CP4','CP6','TP8','P7','P5','P3','P1','PZ','P2','P4','P6','P8','PO7','PO5','PO3','POZ','PO4','PO6','PO8','CB1','O1','OZ','O2','CB2']
import pandas as pd
import matplotlib.pyplot as plt
from mne.time_frequency import psd_welch
labels = [1,0,-1,-1,0,1,-1,0,1,1,0,-1,0,1,-1]
label_meanings = {1:"positive", 0:"neutral", -1:"negative"}
extrapolations = ['box', 'head', 'local']
#fig, axes = plt.subplots(figsize=(8.5, 2.5), ncols=3)
i = 0
fig, ax = plt.subplots(figsize=(8.5, 2.5)) 
 # 创建mne信息对象
sfreq = 1 # 采样频率
seed_montage = mne.channels.read_montage('seed')
ch_types = ['eeg']*62
eeg_info = mne.create_info(ch_names=channels_name,sfreq=sfreq, ch_types=ch_types,montage=seed_montage)

for trail_no in range(1,4,1):
    subject_no = 0
    eegdata = sio.loadmat(subject_names[subject_no])
    trail_data = pd.DataFrame(eegdata[abbreviation[subject_no]+'_eeg'+str(trail_no)], index=channels_name)
    
    del eegdata
    import numpy as np
    from scipy.signal import welch
    f,Pxx_den = welch(trail_data,1.5,nperseg=200)
    Pxx_den = 10*np.log10(Pxx_den)
    Pxx_den /= np.max(Pxx_den)
    f =100*f
    freq_bands = dict(
        delta=(2, 4), theta=(5, 7), alpha=(8, 12), beta=(15, 29), gamma=(30, 45))
    power = {band: np.squeeze(Pxx_den[:,np.where((f >= lf) & (f <= hf))],axis=1) for band, (lf, hf) in freq_bands.items()}
    
    

    for band_name in freq_bands:
        ax = plt.subplot(3,5,i+1)
        band_data = power[band_name]
        if i%5==0:
            ax.set_ylabel(label_meanings[labels[trail_no-1]],fontsize=15)
        
        raw_data =mne.io.RawArray(band_data, eeg_info)
        
        cmap='RdBu_r'
        band_data = np.max(band_data,axis=-1)
        im,cn=mne.viz.plot_topomap(band_data,axes=ax,pos=eeg_info,sensors=False,show=True,names=channels_name,show_names=False)
        mne.viz.topomap._add_colorbar(ax, im, cmap)
        ax.set_title(band_name, fontsize=14)
        
        i += 1

plt.show()
